# -*- coding: utf-8 -*-

import numpy as np
import operator
import os
import matplotlib.pyplot as plt

def knn(k,testdata,traindata,labels):
    #testdata:[feature1,2,3...N]
    #traindata:[[feature1,2,3...N],[],[]]
    traindatasize=traindata.shape[0]
    newtestdata = np.tile(testdata,(traindatasize,1))
    dif=newtestdata-traindata
    sqdif=dif**2
    sumsqdif=sqdif.sum(axis=1)
    distance=sumsqdif**0.5
    sortdistance=distance.argsort()

    count={}
    for i in range(0,k):
        vote=labels[sortdistance[i]]
        count[vote]=count.get(vote,0)+1
    sortcount=sorted(count.items(),key=operator.itemgetter(1),reverse=True)
    return sortcount[0][0]
 
#data load
def datatoarray(fname):
    arr=[]
    fh=open(fname)
    for i in range(0,32):
        thisline = fh.readline()
        for j in range(0,32):
            arr.append(int(thisline[j]))
    return arr
    ##一维数组存储即可，后面一块处理
    
#get label
def seplabel(fname):
    filestr=fname.split(".")[0]
    label = int(filestr.split("_")[0])
    return label

#build train data

def traindata():
    labels = []  # save label
    trainfile = os.listdir("../traindata")
    num = len(trainfile)
    trainarr = np.zeros((num, 1024))  # save traindata, this is ndarray type(from numpy)

    for i in range(0, num):
        thisname = trainfile[i]
        thislabel = seplabel(thisname)
        labels.append(thislabel)
        trainarr[i, :] = datatoarray("../traindata/" + thisname)

    return trainarr, labels

def get_testdata_count():
    testfile = os.listdir("../testdata")
    num = len(testfile)
    return num

def testdata():
    num = get_testdata_count()
    testarr = np.zeros([num,1024])  #num ndarrays, 1024 element of every ndarray
    testlabel = []
    testfile = os.listdir("../testdata")

    for i in range(0, num):
        testarr[i, :] = datatoarray("../testdata/" + testfile[i])
        thislabel = seplabel(testfile[i])
        testlabel.append(thislabel)

    return testarr, testlabel

# main
"""
trainarr, labels = traindata()
testdata = datatoarray("./img/9.txt")
result = knn(3, testdata, trainarr, labels)
print(result)
"""

# Test the algorithm accuracy
trainarr, train_label = traindata()
testarr, test_label = testdata()
testnum = get_testdata_count()
knn_param_range = range(1,10)
right_count = []
wrong_count = []

for n, knn_param in enumerate(knn_param_range):
    right_count.append(0)
    wrong_count.append(0)
    for i in range(testnum):
        knn_result = knn(knn_param, testarr[i], trainarr, train_label)
        if knn_result == test_label[i]:
            #result is right
            right_count[n] = right_count[n] + 1
        else:
            wrong_count[n] = wrong_count[n] + 1

#show the accuracy of different knn_param
for n, knn_param in enumerate(knn_param_range):
    print("knn_param=", knn_param, "right count=", right_count[n], "wrong count=", wrong_count[n])

#save the result
np.savez("./result.npz", f_knn_param_range = knn_param_range, f_testnum = testnum, f_wrong_count = wrong_count)

saved_result = np.load("./result.npz")

plt.plot(saved_result["f_knn_param_range"], saved_result["f_wrong_count"], 'o-')
plt.title('wrong count list')
plt.ylabel('wrong count')
plt.show()

#draw a pic


'''
print("total count = ", testnum)
print("right count = ", right_count)
print("wrong count = ", wrong_count)
'''




